from Parameters import *
from PolymorphismTables import PolymorphismTables
import SiteFrequencySpectrum as sfs
from Polyspect_Operator import *

import scipy.stats.mstats as mstats
from numpy.random import gamma
import infiniteSitesTools as ist
import matplotlib.pyplot as plt
import numpy.matlib as np
from math import sqrt


fig = plt.figure()


p = Population()
p.numChromosomes = 28
p.numSegregatingSites = 0
p.rateRecombination = 0.0
p.sequenceLength = 25000
p.N = 10000

mutationRate = 1E-8

d = Demography(p.N)
simParams = SimulationParameters( p, d )
linOp = ist.getOp_SegregatingSites( p.numChromosomes )

############## Gamma theta 

mutationRates = (gamma(2, size=[1000]) * 0.5 * mutationRate)
allSSList = []
p.numSimulations = 10

for m in mutationRates:
	p.rateMutation = m
	polyOp = Polyspect( linOp, simParams )
	ssList = polyOp.getPolyspectValues()
	print m, ssList
	allSSList.extend(ssList)
	
print allSSList

ax = fig.add_subplot(1,2,1)
ax.hist(allSSList, 40)
ax.set_title( ('Theta from Gamma - variance %1.3f' % (mstats.tvar(allSSList)) ) )


p.numSimulations = 10000

############## Theta 

p.rateMutation = mutationRate
polyOp = Polyspect( linOp, simParams )
ssList = polyOp.getPolyspectValues()


ax = fig.add_subplot(1,2,2)
ax.hist(ssList, 40)
ax.set_title( ('Given theta - variance %1.3f' % (mstats.tvar(ssList)) ) )



############## Seg site

#p.numSegregatingSites = polyOp.getAverageSegSites()
#	
#print simParams
##linOp = ist.getOp_TajimasD( p.numChromosomes )
#polyOp = Polyspect( linOp, simParams )
#ssList = polyOp.getPolyspectValues()
#
#ax = fig.add_subplot(3,1,3)
#ax.hist(ssList)
#ax.set_title( ('Given S - variance %1.3f' % (mstats.tvar(ssList)) ) )



plt.show()